import torch 
import torch.nn as nn 
import matplotlib.pyplot as plt 

class MultiLinear(nn.Module):
    def __init__(self):
        super(MultiLinear, self).__init__()

        self.layers = nn.Sequential(
            nn.Linear(1, 10),  
            nn.Linear(10, 5),
            nn.Sigmoid(),
            nn.Linear(5, 1),
        )
    def forward(self, x):
        pred = self.layers(x)
        return pred 

if __name__ == "__main__":
    
    w = nn.Linear(1, 1)

    t1 = torch.tensor([[1], [2], [3], [4], [5], [6], [7], [8]], dtype=torch.float32)
    target = torch.tensor([[3, 5, 7, 10, 11, 12, 15, 17]])
    l1 = torch.linspace(0, 10, 1000).view(1000, 1)
    t2 = w(l1)
    plt.scatter(t1.view(-1).detach().numpy(), target.detach().numpy(), color="red")
    plt.plot(l1.detach().numpy(), t2.detach().numpy(), color="orange")
    plt.show()

    # print(w(t1))
    w.train()
    optimizer = torch.optim.Adam(w.parameters(), lr=0.01)
    ## 开始训练
    for i in range(10000):
        optimizer.zero_grad()
        pred = w(t1).view(-1)
        loss = ((pred - target)**2).sum() / 8
        print(loss.item())
        if loss.item() < 0.3:
            break 
        loss.backward()
        optimizer.step()

    w.eval()
    print(w(t1))

    l1 = torch.linspace(0, 10, 1000).view(1000, 1)
    print(l1.shape)
    t2 = w(l1)
    plt.scatter(t1.view(-1).detach().numpy(), target.detach().numpy(), color="red")
    plt.plot(l1.detach().numpy(), t2.detach().numpy())
    plt.show()

    model = MultiLinear()
    t1 = torch.tensor([[1], [2], [3], [4], [5], [6], [7], [8]], dtype=torch.float32)
    target = torch.tensor([[3, 5, 7, 10, 11, 12, 15, 17]])
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    ## 开始训练
    for i in range(20000):
        optimizer.zero_grad()
        pred = model(t1).view(-1)
        loss = ((pred - target)**2).sum() / 8
        print(loss.item())
        if loss.item() < 0.04:
            break 
        loss.backward()
        optimizer.step()

    model.eval()
    print(model(t1))

    l1 = torch.linspace(0, 10, 1000).view(1000, 1)
    # print(l1.shape)
    t2 = model(l1)
    plt.scatter(t1.view(-1).detach().numpy(), target.detach().numpy(), color="red")
    plt.plot(l1.detach().numpy(), t2.detach().numpy())
    plt.show()
